import os
import numpy as np
from PIL import Image

# 定义标签信息
label_info = {
    1: (128, 0, 0),  # building - Red
    2: (128, 128, 128),  # sky - Grey
    3: (128, 64, 128),  # road - Pink
    4: (128, 128, 0),  # vegetation - Dark yellow
    5: (0, 0, 192),  # sidewalk/pave - Blue
    6: (64, 0, 128),  # car - Purple
    7: (64, 64, 0),  # pedestrian - Yellow-brown
    8: (0, 128, 192),  # cyclist - Light blue
    9: (192, 128, 128),  # signage - Salmon
    10: (64, 64, 128),  # fence/wall - Grey-purple
    11: (192, 192, 128),  # pole - Light yellow
    0: (0, 0, 0)  # other - Black
}


def visualize_bin_to_image(bin_file, output_image_path, width, height, num_labels):
    with open(bin_file, 'rb') as f:
        data = np.fromfile(f, dtype=np.float32)
    data = data.reshape((height, width, num_labels))

    # 创建图像对象
    img = Image.new('RGB', (width, height))

    # 设置每个像素的颜色
    for y in range(height):
        for x in range(width):
            label_probabilities = data[y, x]
            label_idx = np.argmax(label_probabilities)  # 获取最大概率对应的标签索引
            color = label_info[label_idx + 1]
            img.putpixel((x, y), color)

    # 保存图像
    img.save(output_image_path)


def batch_visualize_bin_to_images(bin_dir, output_dir, width, height, num_labels):
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    bin_files = [f for f in os.listdir(bin_dir) if f.endswith('.bin')]

    for bin_file in bin_files:
        bin_file_path = os.path.join(bin_dir, bin_file)
        output_file_name = f"{os.path.splitext(bin_file)[0]}_color.png"
        output_image_path = os.path.join(output_dir, output_file_name)
        visualize_bin_to_image(bin_file_path, output_image_path, width, height, num_labels)
        print(f"Processed: {bin_file} -> {output_image_path}")


# 示例用法
data_root = 'E:\change_data\c_data_0905key2'
# bin_file_path = os.path.join(data_root, 'label_2d_bin', '000008.bin')
# output_image_path = os.path.join(data_root, 'test', 'example.png')
bin_dir = os.path.join(data_root, 'label_2d_bin')  # 批量护理用法
output_dir = os.path.join(data_root, 'label_visual')
width = 640  # 图像宽度
height = 480  # 图像高度
num_labels = 11  # 标签数量
# visualize_bin_to_image(bin_file_path, output_image_path, width, height, num_labels)
batch_visualize_bin_to_images(bin_dir, output_dir, width, height, num_labels)  # 批量处理用法

'''19类的bin文件转化'''
# import os
# import numpy as np
# from PIL import Image
# '''将映射后的灰度图片转化为彩色图片，用于检查映射是否正确'''
# # 定义标签信息
# label_info = {
#     1: (244, 35, 232),  # building - Red
#     2: (70, 70, 70),  # sky - Grey
#     3: (102, 102, 156),  # road - Pink
#     4: (190, 153, 153),  # vegetation - Dark yellow
#     5: (153, 153, 153),  # sidewalk/pave - Blue
#     6: (250, 170, 30),  # car - Purple
#     7: (220, 220, 0),  # pedestrian - Yellow-brown
#     8: (107, 142, 35),  # cyclist - Light blue
#     9: (152, 251, 152),  # signage - Salmon
#     10: (70, 130, 180),  # fence/wall - Grey-purple
#     11: (220, 20, 60),  # pole - Light yellow
#     12: (255, 0, 0),  # traffic light - Red
#     13: (0, 0, 142),  # traffic sign - Blue
#     14: (0, 0, 70),  # vegetation - Dark yellow
#     15: (0, 60, 100),  # building - Red
#     16: (0, 80, 100),  # signage - Salmon
#     17: (0, 0, 230),  # cyclist - Light blue
#     18: (119, 11, 32),  # sky - Grey,  # pedestrian - Yellow-brown
#     19: (0, 0, 0),  # other - Blacks
#     0: (128, 64, 128)  # other - Black
# }
#
#
# def visualize_bin_to_image(bin_file, output_image_path, width, height, num_labels):
#     with open(bin_file, 'rb') as f:
#         data = np.fromfile(f, dtype=np.float32)
#
#     data = data.reshape((height, width, num_labels))
#
#     # 创建图像对象
#     img = Image.new('RGB', (width, height))
#
#     # 设置每个像素的颜色
#     for y in range(height):
#         for x in range(width):
#             label_probabilities = data[y, x]
#             label_idx = np.argmax(label_probabilities)  # 获取最大概率对应的标签索引
#             color = label_info[label_idx]
#             img.putpixel((x, y), color)
#
#     # 保存图像
#     img.save(output_image_path)
#
#
# # 示例用法
# data_root = 'D:\guomengqi\毕业倒计时\课题\语义地图处理部分\deeplab_v3\mapping_img_change\data_change0528'
# bin_file_path = os.path.join(data_root, 'label_2d_bin_raw', '000005.bin')
# output_image_path = os.path.join(data_root, 'test', 'example.png')
# width = 640  # 图像宽度
# height = 480  # 图像高度
# num_labels = 19  # 标签数量
# visualize_bin_to_image(bin_file_path, output_image_path, width, height, num_labels)
